library(ggpubr)
library(dplyr)
library(tidyr)
library(ggplot2)
library(cowplot)
library(RColorBrewer)
library(viridis)


if (!file.exists("output")) {
  dir.create("output")
}

# load data

source("HNC_metadata_tidy.R")
source("data_handling.R")

sbsCOSMIC <- read.csv("Input_signature_attributions/Input_SBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")
dbsCOSMIC <- read.csv("Input_signature_attributions/Input_DBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")
idCOSMIC <- read.csv("Input_signature_attributions/Input_ID_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")

additional_df <- list(sbsCOSMIC, dbsCOSMIC, idCOSMIC)

cosmic <- additional_df[[1]]
for (n in 2:length(additional_df)) {
  cosmic <- cosmic %>% left_join(additional_df[[n]], by = "donor_id")
}

# merge with cosmic sigantures
# Samples with a relative SBS7a-c burden of >10% are categorized as positive for UV exposure
data <- data %>%
  left_join(cosmic, by = "donor_id") %>%
  rename(anatomic_site = icd_description) %>%
  mutate(
    `SBS7a-c` = SBS7a + SBS7b + SBS7c,
    SBS7ac_rel = `SBS7a-c` / Mutational.burden,
    `SBS7a-c_cat` = ifelse(SBS7ac_rel > 0.10, 1, 0)
  )

data <- dichotomizeSigs(metadata = data, sigsn = cosmic)

data %>%
  # filter samples with UV
  filter(`SBS7a-c_cat` == 1) %>%
  # simplify anatomic site and spread data frame
  mutate(anatomic_site = case_when(
    anatomic_site == "Malignant neoplasm of lip" ~ "Lip",
    anatomic_site == "Malignant neoplasm of other and unspecified parts of tongue" ~ "Tongue",
    anatomic_site == "Malignant neoplasm of floor of mouth" ~ "Floor of mouth",
    TRUE ~ anatomic_site
  )) %>%
  select("donor_id", anatomic_site, all_of(names(cosmic[-1]))) %>%
  gather(Signature, value, -c(donor_id, anatomic_site)) %>%
  na.omit() %>%
  mutate(
    sigtype = case_when(grepl("SBS", Signature) ~ "SBS", grepl("DBS", Signature) ~ "DBS", grepl("ID", Signature) ~ "ID"),
    sigtype = factor(sigtype, level = c("SBS", "DBS", "ID")),
    Signature = ifelse(Signature %in% c("SBS7a", "SBS7b", "SBS7c", "DBS1", "ID13"), Signature, "others"),
    Signature = factor(Signature, levels = c("others", "SBS7a", "SBS7b", "SBS7c", "DBS1", "ID13")),
    anatomic_site = factor(anatomic_site, levels = c("Lip", "Tongue", "Floor of mouth"))
  ) %>%
  group_by(donor_id, anatomic_site, sigtype, Signature) %>%
  summarize(value = sum(value)) %>%
  # plot
  ggplot(aes(x = donor_id, y = value, fill = Signature)) +
  facet_grid(sigtype ~ anatomic_site, scales = "free", space = "free_x", switch = "y") +
  geom_col() +
  scale_y_continuous(position = "right") +
  scale_fill_manual(values = c("grey", "#0c2c84", "#225ea8", "#1d91c0", "#41b6c4", "#7fcdbb")) +
  ylab("Mutation burden") +
  xlab("") +
  ggtitle("") +
  theme_bw() +
  theme(
    plot.title = element_text(face = "bold", size = 12),
    axis.text.x = element_blank(),
    axis.title.y = element_text(size = 12),
    strip.text = element_text(face = "bold"),
    panel.grid = element_blank(),
    strip.background = element_rect(fill = "white"),
    legend.text = element_text(size = 12),
    legend.title = element_blank(),
    legend.position = "bottom"
  )

ggsave(paste0("output/Figure_6b_", Sys.Date(), ".pdf"), device = "pdf", width = 6, height = 4, dpi = 700)
